{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2017 The TensorFlow Authors. All Rights Reserved.  \n",
    "  \n",
    " Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    " you may not use this file except in compliance with the License.  \n",
    " You may obtain a copy of the License at  \n",
    "  \n",
    "     http://www.apache.org/licenses/LICENSE-2.0  \n",
    "  \n",
    " Unless required by applicable law or agreed to in writing, software  \n",
    " distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    " See the License for the specific language governing permissions and  \n",
    " limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only for when running on Colab:\n",
    "import sys\n",
    "if 'google.colab' in sys.modules:\n",
    "    # Get the dependency .py files, if any.\n",
    "    ! git clone https://github.com/GoogleCloudPlatform/cloudml-samples.git\n",
    "    ! cp cloudml-samples/census/keras/trainer/* .\n",
    "\n",
    "    # Authenticate the user for better GCS access.\n",
    "    # Copy verification code into the text field to continue.\n",
    "    from google.colab import auth\n",
    "    auth.authenticate_user()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"This code implements a Feed forward neural network using Keras API.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import glob\n",
    "import os\n",
    "from keras.callbacks import Callback\n",
    "from keras.callbacks import ModelCheckpoint\n",
    "from keras.callbacks import TensorBoard\n",
    "from keras.models import load_model\n",
    "from tensorflow.python.lib.io import file_io\n",
    "import model as model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_SIZE = 55\n",
    "CLASS_SIZE = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CHUNK_SIZE specifies the number of lines  \n",
    " to read in case the file is very large"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CHUNK_SIZE = 5000\n",
    "CHECKPOINT_FILE_PATH = 'checkpoint.{epoch:02d}.hdf5'\n",
    "CENSUS_MODEL = 'census.hdf5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContinuousEval(Callback):\n",
    "  \"\"\"Continuous eval callback to evaluate the checkpoint once\n",
    "\n",
    "     every so many epochs.\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self,\n",
    "               eval_frequency,\n",
    "               eval_files,\n",
    "               learning_rate,\n",
    "               job_dir,\n",
    "               steps=1000):\n",
    "    self.eval_files = eval_files\n",
    "    self.eval_frequency = eval_frequency\n",
    "    self.learning_rate = learning_rate\n",
    "    self.job_dir = job_dir\n",
    "    self.steps = steps\n",
    "\n",
    "  def on_epoch_begin(self, epoch, logs={}):\n",
    "    \"\"\"Compile and save model.\"\"\"\n",
    "    if epoch > 0 and epoch % self.eval_frequency == 0:\n",
    "      # Unhappy hack to work around h5py not being able to write to GCS.\n",
    "      # Force snapshots and saves to local filesystem, then copy them over to GCS.\n",
    "      model_path_glob = 'checkpoint.*'\n",
    "      if not self.job_dir.startswith('gs://'):\n",
    "        model_path_glob = os.path.join(self.job_dir, model_path_glob)\n",
    "      checkpoints = glob.glob(model_path_glob)\n",
    "      if len(checkpoints) > 0:\n",
    "        checkpoints.sort()\n",
    "        census_model = load_model(checkpoints[-1])\n",
    "        census_model = model.compile_model(census_model, self.learning_rate)\n",
    "        loss, acc = census_model.evaluate_generator(\n",
    "            model.generator_input(self.eval_files, chunk_size=CHUNK_SIZE),\n",
    "            steps=self.steps)\n",
    "        print('\\nEvaluation epoch[{}] metrics[{:.2f}, {:.2f}] {}'.format(\n",
    "            epoch, loss, acc, census_model.metrics_names))\n",
    "        if self.job_dir.startswith('gs://'):\n",
    "          copy_file_to_gcs(self.job_dir, checkpoints[-1])\n",
    "      else:\n",
    "        print('\\nEvaluation epoch[{}] (no checkpoints found)'.format(epoch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_evaluate(args):\n",
    "  census_model = model.model_fn(INPUT_SIZE, CLASS_SIZE)\n",
    "  try:\n",
    "    os.makedirs(args.job_dir)\n",
    "  except:\n",
    "    pass\n",
    "\n",
    "  # Unhappy hack to workaround h5py not being able to write to GCS.\n",
    "  # Force snapshots and saves to local filesystem, then copy them over to GCS.\n",
    "  checkpoint_path = CHECKPOINT_FILE_PATH\n",
    "  if not args.job_dir.startswith('gs://'):\n",
    "    checkpoint_path = os.path.join(args.job_dir, checkpoint_path)\n",
    "\n",
    "  # Model checkpoint callback.\n",
    "  checkpoint = ModelCheckpoint(\n",
    "      checkpoint_path,\n",
    "      monitor='val_loss',\n",
    "      verbose=1,\n",
    "      period=args.checkpoint_epochs,\n",
    "      mode='min')\n",
    "\n",
    "  # Continuous eval callback.\n",
    "  evaluation = ContinuousEval(args.eval_frequency, args.eval_files,\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\targs.learning_rate, args.job_dir)\n",
    "\n",
    "  # Tensorboard logs callback.\n",
    "  tb_log = TensorBoard(\n",
    "      log_dir=os.path.join(args.job_dir, 'logs'),\n",
    "      histogram_freq=0,\n",
    "      write_graph=True,\n",
    "      embeddings_freq=0)\n",
    "\n",
    "  callbacks = [checkpoint, evaluation, tb_log]\n",
    "\n",
    "  census_model.fit_generator(\n",
    "      model.generator_input(args.train_files, chunk_size=CHUNK_SIZE),\n",
    "      steps_per_epoch=args.train_steps,\n",
    "      epochs=args.num_epochs,\n",
    "      callbacks=callbacks)\n",
    "\n",
    "  # Unhappy hack to workaround h5py not being able to write to GCS.\n",
    "  # Force snapshots and saves to local filesystem, then copy them over to GCS.\n",
    "  if args.job_dir.startswith('gs://'):\n",
    "    census_model.save(CENSUS_MODEL)\n",
    "    copy_file_to_gcs(args.job_dir, CENSUS_MODEL)\n",
    "  else:\n",
    "    census_model.save(os.path.join(args.job_dir, CENSUS_MODEL))\n",
    "\n",
    "  # Convert the Keras model to TensorFlow SavedModel.\n",
    "  model.to_savedmodel(census_model, os.path.join(args.job_dir, 'export'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "h5py workaround: copy local models over to GCS if the job_dir is GCS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def copy_file_to_gcs(job_dir, file_path):\n",
    "  with file_io.FileIO(file_path, mode='rb') as input_f:\n",
    "    with file_io.FileIO(\n",
    "        os.path.join(job_dir, file_path), mode='w+') as output_f:\n",
    "      output_f.write(input_f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\n",
    "    '--train-files',\n",
    "    nargs='+',\n",
    "    help='Training file local or GCS',\n",
    "    default='gs://cloud-samples-data/ml-engine/census/data/adult.data.csv')\n",
    "parser.add_argument(\n",
    "    '--eval-files',\n",
    "    nargs='+',\n",
    "    help='Evaluation file local or GCS',\n",
    "    default='gs://cloud-samples-data/ml-engine/census/data/adult.test.csv')\n",
    "parser.add_argument(\n",
    "    '--job-dir',\n",
    "    type=str,\n",
    "    help='GCS or local dir to write checkpoints and export model',\n",
    "    default='/tmp/census-keras')\n",
    "parser.add_argument(\n",
    "    '--train-steps',\n",
    "    type=int,\n",
    "    default=100,\n",
    "    help=\"\"\"\\\n",
    "      Maximum number of training steps to perform\n",
    "      Training steps are in the units of training-batch-size.\n",
    "      So if train-steps is 500 and train-batch-size if 100 then\n",
    "      at most 500 * 100 training instances will be used to train.\"\"\")\n",
    "parser.add_argument(\n",
    "    '--eval-steps',\n",
    "    help='Number of steps to run evalution for at each checkpoint',\n",
    "    default=100,\n",
    "    type=int)\n",
    "parser.add_argument(\n",
    "    '--train-batch-size',\n",
    "    type=int,\n",
    "    default=40,\n",
    "    help='Batch size for training steps')\n",
    "parser.add_argument(\n",
    "    '--eval-batch-size',\n",
    "    type=int,\n",
    "    default=40,\n",
    "    help='Batch size for evaluation steps')\n",
    "parser.add_argument(\n",
    "    '--learning-rate',\n",
    "    type=float,\n",
    "    default=0.003,\n",
    "    help='Learning rate for SGD')\n",
    "parser.add_argument(\n",
    "    '--eval-frequency',\n",
    "    default=10,\n",
    "    help='Perform one evaluation per n epochs')\n",
    "parser.add_argument(\n",
    "    '--first-layer-size',\n",
    "    type=int,\n",
    "    default=256,\n",
    "    help='Number of nodes in the first layer of DNN')\n",
    "parser.add_argument(\n",
    "    '--num-layers',\n",
    "    type=int,\n",
    "    default=2,\n",
    "    help='Number of layers in DNN')\n",
    "parser.add_argument(\n",
    "    '--scale-factor',\n",
    "    type=float,\n",
    "    default=0.25,\n",
    "    help=\"\"\"Rate of decay size of layer for Deep Neural Net.\n",
    "      max(2, int(first_layer_size * scale_factor**i))\"\"\")\n",
    "parser.add_argument(\n",
    "    '--eval-num-epochs',\n",
    "    type=int,\n",
    "    default=1,\n",
    "    help='Number of epochs during evaluation')\n",
    "parser.add_argument(\n",
    "    '--num-epochs',\n",
    "    type=int,\n",
    "    default=20,\n",
    "    help='Maximum number of epochs on which to train')\n",
    "parser.add_argument(\n",
    "    '--checkpoint-epochs',\n",
    "    type=int,\n",
    "    default=5,\n",
    "    help='Checkpoint per n training epochs')\n",
    "\n",
    "args, _ = parser.parse_known_args()\n",
    "train_and_evaluate(args)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
